import gymnasium as gym
import ale_py

class RewardPenaltyWrapper(gym.Wrapper):
    def __init__(self, env, frame_penalty=-0.1, life_loss_penalty=-10):
        super(RewardPenaltyWrapper, self).__init__(env)
        self.frame_penalty = frame_penalty
        self.life_loss_penalty = life_loss_penalty
        self.previous_lives = 0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info.get('lives', 0)  # 初始生命值
        return obs, info

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)

        # if reward != 0:
        #     reward /= 10.0 # 缩放奖励
        
        # 处理生命减少时的惩罚
        current_lives = info.get('lives', self.previous_lives)
        if current_lives < self.previous_lives:
            reward += self.life_loss_penalty
            self.previous_lives = current_lives
        
        return obs, reward, done, truncated, info

gym.register_envs(ale_py)

# Initialize the environment
env = RewardPenaltyWrapper(gym.make('ALE/BasicMath-v5', render_mode="human"))
# env = gym.make('ALE/BankHeist-v5', render_mode="human")

# Reset the environment to get the initial state
state = env.reset()
total_reward = 0
# Run a loop to play the game
for _ in range(10000):
    # Take a random action
    action = env.action_space.sample()

    # Get the next state, reward, done flag, and info from the environment
    state, reward, done, trunc, info = env.step(action)
    total_reward += reward
    if reward != 0:
        print("reward: ", reward)
        print("info: ", info)

    # If done, reset the environment
    if done or trunc:
        break

print("Total reward: ", total_reward)

# Close the environment
env.close()